/**
 * @file elimination.h
 * 
 * @brief Gauss elimination... etc
 */
#ifndef _ELIMINATION_H_
#define _ELIMINATION_H_

#include "core.h"
#include <cmath>

NAMESPACE_LA_BEGIN

template <typename T>
constexpr bool equals_zero(const T& v)
{
    return v == 0;
}

constexpr float zero_threshold_f = 1e-8;
constexpr double zero_threshold_d = 1e-10;
constexpr long double zero_threshold_ld = 1e-12;

template <>
constexpr bool equals_zero<float>(const float& v)
{
    return std::abs(v) < zero_threshold_f;
}

template <>
constexpr bool equals_zero<double>(const double& v)
{
    return std::abs(v) < zero_threshold_d;
}

template <>
constexpr bool equals_zero<long double>(const long double& v)
{
    return std::abs(v) < zero_threshold_ld;
}

/**
 * @function eliminate
 * 
 * @brief underlying implementation of elimination
 */
template <typename T>
int eliminate(Matrix<T>& mat, Matrix<T>& aug, int pivot_cols[])
{
    int n_rows = mat.m(), n_cols = mat.n();
    int rank = 0, j = 0; // j for column position

    /**
     * make sure pivot_cols is big enough to hold M integers ！
     * pivot column index of each row, -1 for no pivot
     */
    std::fill_n(pivot_cols, n_rows, -1);

    for (int i = 0; i < n_rows; ++i) {
        // iterate the rows
        while (j < n_cols) {
            bool pivot_found = false;
            int max_index = i;

            for (int k = i; k < n_rows; ++k) {
                if (equals_zero(mat(k, j)))
                    continue; // A[k, j] cannot be chosen as a pivot
                pivot_found = true;
                if (std::abs(mat(max_index, j)) < std::abs(mat(k, j))) {
                    max_index = k;
                    /**
                     * @note: there's much difference between abs & std::abs.
                     * abs is a C library routine declared in stdlib.h which only performs on
                     * integer types. (int only, to be specific)
                     * in contrast, std::abs performs on integers and floating point numbers.
                     * its return type will be double if an integer is given as argument.
                     */
                }
            }

            if (pivot_found) {
                // we find a pivot, increase rank by 1
                ++rank;
                // adjust the rows, swap the row with the maximum pivot & current row
                if (i != max_index) { // no need to swap if the row is already here
                    for (int k = j; k < n_cols; ++k)
                        std::swap(mat(i, k), mat(max_index, k));
                    // also operate on the additional matrix
                    for (int k = 0; k < aug.n(); ++k)
                        std::swap(aug(i, k), aug(max_index, i));
                }
                break;
            }

            ++j; // pivot not found, next column
        }

        if (j == n_cols) // no more pivot can be found , stop.
            break;
        
        pivot_cols[i] = j; // record it !
        
        // elimination step
        for (int k = i + 1; k < n_rows; ++k) {
            T multiplier = mat(k, j) / mat(i, j); // pivot: A[i, j]
            for (int l = j; l < n_cols; ++l)
                mat(k, l) -= mat(i, l) * multiplier;
            for (int l = 0; l < aug.n(); ++l)
                aug(k, l) -= aug(i, l) * multiplier;
        }
        ++j; // next column
    }
    return rank;
}

template <typename T>
void back_substitute(Matrix<T>& mat, Matrix<T>& aug, const int pivot_cols[])
{
    for (int i = mat.m() - 1; i >= 0; --i) {
        //std::cout << "dgf" << i << std::endl;
        if (pivot_cols[i] == -1)
            continue; // no pivot
        // divides pivot to one
        T multiplier = T(1) / mat(i, pivot_cols[i]);
        for (int j = pivot_cols[i]; j < mat.n(); ++j)
            mat(i, j) *= multiplier;
        for (int j = 0; j < aug.n(); ++j)
            aug(i, j) *= multiplier;
        for (int k = i - 1; k >= 0; --k) {
            multiplier = mat(k, pivot_cols[i]);
            for (int j = pivot_cols[i]; j < mat.n(); ++j)
                mat(k, j) -= mat(i, j) * multiplier;
            for (int j = 0; j < aug.n(); ++j)
                aug(k, j) -= aug(i, j) * multiplier;
        }
    }
}

/**
 * @function eliminate
 * 
 * @brief performs in situ elimination --> row echelon form
 *      second version. performs on an augmented matrix ([mat | aug]).
 * @return rank of the matrix
 */
template <typename T>
int eliminate(Matrix<T>& mat, Matrix<T>& aug)
{
    assert(mat.m() > 0 && mat.n() > 0 && mat.m() == aug.m());
    int *pivot_cols = new int[mat.m()];
    int rank = eliminate(mat, aug, pivot_cols);
    delete[] pivot_cols;
    return rank;
}

/**
 * @function eliminate
 * 
 * @brief performs in situ elimination --> row echelon form
 *      second version. performs on a matrix and a vector.
 * @return rank of the matrix
 */
template <typename T>
int eliminate(Matrix<T>& mat, Vector<T>& vec)
{
    assert(mat.m() > 0 && mat.n() > 0 && mat.m() == vec.n());
    int *pivot_cols = new int[mat.m()];
    Matrix<T> aug(vec.n(), 1, vec.view());
    int rank = eliminate(mat, aug, pivot_cols);
    delete[] pivot_cols;
    return rank;
}

/**
 * @function eliminate
 * 
 * @brief performs in situ elimination --> row echelon form
 * @return rank of the matrix
 */
template <typename T>
int eliminate(Matrix<T>& mat)
{
    assert(mat.m() > 0 && mat.n() > 0);
    int *pivot_cols = new int[mat.m()];
    Matrix<T> aug; // an empty matrix
    int rank = eliminate(mat, aug, pivot_cols);
    delete[] pivot_cols;
    return rank;    
}

/**
 * @function to_rref
 * 
 * @brief reduce a matrix into its RREF (Reduced Row Echelon Form) 
 * @return rank of the matrix
 */
template <typename T>
int to_rref(Matrix<T>& mat)
{
    assert(mat.m() > 0 && mat.n() > 0);
    int *pivot_cols = new int[mat.m()];
    Matrix<T> aug; // an empty matrix
    int rank = eliminate(mat, aug, pivot_cols);
    back_substitute(mat, aug, pivot_cols);
    delete[] pivot_cols;
    return rank;    
}

NAMESPACE_LA_END

#endif // _ELIMINATION_H_
